from enum import Enum

from tortoise import fields, models


class StudyPlanStatus(str, Enum):
    FORECASTED = "forecasted"
    APPROVED = "approved"
    COMPLETED = "completed"


class StudyPlan(models.Model):
    year = fields.IntField()
    month = fields.IntField()
    study: fields.ForeignKeyRelation["models.MedicalStudy"] = fields.ForeignKeyField("models.MedicalStudy")
    number_of_studies = fields.IntField()
    status = fields.CharEnumField(StudyPlanStatus, default=StudyPlanStatus.FORECASTED)

    class Meta:
        unique_together = (("year", "month", "study"),)
